#include <iostream>
#include <fst/queue.h>
#include <fst/vector-fst.h>
#include <fst/symbol-table.h>
#include <fst/arcsort.h>
#include <fst/matcher-fst.h>
#include <fst/compose.h>
#include <fst/rmepsilon.h>
#include <fst/determinize.h>
#include <fst/minimize.h>
#include <fstx/all-paths.h>

using namespace std;
using namespace fst;



template<class Weight> 
struct LevensteianFstOpts {
	LevensteianFstOpts (): del_cost(Weight::One()),
		in_cost(Weight::One()), sub_cost(Weight::One()) { };
	Weight del_cost;
	Weight in_cost;
	Weight sub_cost;
};

int FindState(int i, int e, int l) {
  return e * l  + i;
}


///Dist is the maximum number of edits
template<class Arc>
void ConstructLevestein(const string& word, int dist, VectorFst<Arc>* ofst, 
	SymbolTable* syms, LevensteianFstOpts<typename Arc::Weight> opts) {
  
	typedef typename Arc::StateId S;
  typedef typename Arc::Label L;
  typedef typename Arc::Weight W;

  S num_states = (word.size() + 1) * (dist + 1);  
  ofst->DeleteStates();
	VectorFst<Arc> fst;
  for (S s = 0; s != num_states; ++s)
    fst.AddState();
  fst.SetStart(0);
  SymbolTable ssyms("ssyms");

  int l = word.size() + 1;
  for (int i = 0; i != word.size(); ++i) {
    stringstream ss;
    ss <<  word[i];
    string c = ss.str();    
    for (int e = 0; e <= dist ; ++e) {
      //match
      S s = FindState(i, e, l);
      S d = FindState(i + 1, e, l);
      L a = syms->AddSymbol(c.c_str());
      fst.AddArc(s, Arc(a, a, W::One(), d));
      if (e < dist) {
        //insert
        s = FindState(i, e, l);
        d = FindState(i, e + 1, l);
				fst.AddArc(s, Arc(1, 1, opts.in_cost, d));
        //delete
        s = FindState(i, e, l);
        d = FindState(i + 1, e + 1, l);
				fst.AddArc(s, Arc(0, 0, opts.del_cost, d));
        //sub
        s = FindState(i, e, l);
        d = FindState(i + 1, e + 1, l);
				fst.AddArc(s, Arc(1, 1, opts.sub_cost, d));
      }    
    }    
  }  
  for (int e = 0; e <= dist; ++e) {
    if (e < dist) {
      S s = FindState(word.size(), e, l);
      S d = FindState(word.size(), e + 1, l);
      fst.AddArc(s, Arc(1, 1, W(1.0f), d)); //W::One()
    }    
    fst.SetFinal(FindState(word.size(), e, l), W::One());
  }
  RmEpsilon(&fst);	
  Determinize(fst, ofst);
  Minimize(ofst);	
  ArcSort(ofst, OLabelCompare<Arc>());
}

template<class S, class W, class L>
struct Element {  
  Element* prev;
  L l;
  S lq;
  S dq;
	W cost;	

	Element() : prev(0), lq(kNoStateId), dq(kNoStateId), cost(W::One()) { }

	Element(Element* prev, L l, S lq, S dq,	W cost) :
		prev(prev), l(l), lq(lq), dq(dq), cost(cost) { }

	~Element() { }

  Element* Extend(L l, S lq, S dq, const W& w) {
    return new Element(this, l, lq, dq, Times(cost, w));    
  }

	void TraceBack(stringstream* ss, const SymbolTable& syms) {
		(*ss) << syms.Find(l);
    if (prev != 0)
      prev->TraceBack(ss, syms);
  }
};


template<class Arc, class E, class M>
void SearchR(E& prev, const Fst<Arc>& levfst, const Fst<Arc>& dictfst,
	const SymbolTable& syms) {
	typedef typename Arc::Weight W;	
	typedef Fst<Arc> FST;	
	
	M m(levfst, fst::MATCH_OUTPUT, 1);	
	if (levfst.Final(prev.lq) != W::Zero() &&
			dictfst.Final(prev.dq) != W::Zero()) {
		stringstream ss;
		prev.TraceBack(&ss, syms);		
		string str = ss.str();
		reverse(str.begin(), str.end());
		LOG(INFO) << str << " " << Times(prev.cost, levfst.Final(prev.lq));
	}

	int extended = 0;
	m.SetState(prev.lq);
	for (ArcIterator<FST> aiter(dictfst, prev.dq); !aiter.Done(); aiter.Next()) {
		const Arc& darc = aiter.Value();
		if (m.Find(darc.ilabel)) {
			while (!m.Done()) {
				const fst::StdArc& larc = m.Value();				
				E e(&prev, darc.ilabel, larc.nextstate, darc.nextstate,
					Times(prev.cost, larc.weight));
				SearchR<Arc, E, M>(e, levfst, dictfst, syms);
				m.Next();
			}
		}
	}
}

template<class Arc>
void SearchR(const Fst<Arc>& levfst, const Fst<Arc>& dictfst,
	const SymbolTable& syms) {
		typedef typename Arc::Label L;
		typedef typename Arc::StateId S;
		typedef typename Arc::Weight W;	
		typedef Fst<Arc> FST;	
		typedef Element<S, W, L> E;
		typedef RhoMatcher< SortedMatcher<FST> > M;
		E e;
		e.prev = 0;
		e.lq = levfst.Start();
		e.dq = dictfst.Start();
		SearchR<Arc, E, M>(e, levfst, dictfst, syms);
}

template<class Arc>
void Search(const Fst<Arc>& levfst, const Fst<Arc>& dictfst, const SymbolTable& syms) {
  typedef typename Arc::Label Label;
  typedef typename Arc::StateId StateId;
  typedef typename Arc::Weight Weight;  

  typedef Fst<Arc> FST;
	typedef RhoMatcher< SortedMatcher<FST> > M;
  M m(levfst, fst::MATCH_OUTPUT, 1);

	typedef Element<StateId, Weight, Label> E;
	LifoQueue<E*> Q;
	vector<E*> elements;	
	E root(0, kNoLabel, levfst.Start(), dictfst.Start(), Weight::One());
  
	Q.Enqueue(&root);
  while (!Q.Empty()) {
    E* e = Q.Head();
    Q.Dequeue();
    if (levfst.Final(e->lq) != Weight::Zero() && dictfst.Final(e->dq) != Weight::Zero()) {
			stringstream ss;
			e->TraceBack(&ss, syms);
			string str = ss.str();
      reverse(str.begin(), str.end());
			LOG(INFO) << str << " " << Times(e->cost, levfst.Final(e->lq));
    }

    int extended = 0;
    m.SetState(e->lq);
    for (ArcIterator<FST> aiter(dictfst, e->dq); !aiter.Done(); aiter.Next()) {
      const Arc& darc = aiter.Value();
      if (m.Find(darc.ilabel)) {
        while (!m.Done()) {
          const fst::StdArc& larc = m.Value();          
					E* f = e->Extend(darc.ilabel, larc.nextstate, darc.nextstate, larc.weight);
					elements.push_back(f);
          Q.Enqueue(f);
          extended++;
          m.Next();
        }
      }
    }       
  }

	for (vector<E*>::iterator it = elements.begin(); it != elements.end(); ++it)
		delete *it;
}

int main(int argc, char** argv) {
	typedef StdArc Arc;
	typedef StdArc::Label L;
	typedef StdArc::StateId S;
	typedef StdArc::Weight W;
	typedef Fst<Arc> FST;
	typedef RhoMatcher< SortedMatcher<FST> > M;
	//typedef SigmaMatcher< SortedMatcher<FST> > M;
	typedef ComposeFstOptions<Arc, M> COpts;
	
	string dictpath(argv[1]);
	StdFst* dictfst = StdFst::Read(dictpath);
	SymbolTable syms(*dictfst->InputSymbols());
	string input(argv[2]);
	int dist = atoi(argv[3]);
	L all_label = syms.Find("<all>");

	StdVectorFst levfst;
	LevensteianFstOpts<W> levopts;
	ConstructLevestein(input, dist, &levfst, &syms, levopts);
	levfst.SetInputSymbols(&syms);
	levfst.SetOutputSymbols(&syms);
	Search(levfst, *dictfst, syms);

	//levfst.Write(root + "lev.ofst");	
	
	//SearchR(levfst, *dictfst, syms);
	
	
	if (true) {
		typedef ComposeFstOptions<Arc, M> COpts;	
		COpts opts;		
		opts.matcher1 = new M(levfst, fst::MATCH_OUTPUT, 1);
		opts.matcher2 = new M(*dictfst, fst::MATCH_NONE);  	

		StdVectorFst ofst(StdComposeFst(levfst, *dictfst, opts));	
		Connect(&ofst);
		//ofst.Write(root + "ld.ofst");
		vector<string> strings;
		AllPaths(ofst, &syms, &strings);
		for (vector<string>::iterator it = strings.begin();
				it != strings.end(); ++it) {
			cerr << *it << endl;
		}
	}
	return 0;
}